#include "UdpServer.h"

using namespace clusnet::net;

UdpServer::UdpServer(char const* host, unsigned short port)
    : UdpServer(boost::asio::ip::udp::endpoint(boost::asio::ip::make_address(host), port))
{
}

UdpServer::UdpServer(unsigned short port, boost::asio::ip::udp::resolver::protocol_type const& proto)
    : UdpServer(boost::asio::ip::udp::endpoint(proto, port))
{
}

UdpServer::UdpServer(boost::asio::ip::udp::endpoint const& endpoint)
    : io_(),
      strand_(io_),
      localEndpoint_(endpoint),
      socket_(io_, localEndpoint_),
      stop_(false),
      messageCallback_(nullptr),
      writeCallback_(nullptr)
{
    struct sigaction sa;
    sa.sa_handler = SIG_IGN;
    sigaction(SIGPIPE, &sa, nullptr);
}

UdpServer::~UdpServer()
{
    assert(threads_.empty());
    assert(sessions_.empty());
    if (!stop_)
    {
        socket_.close();
    }
}

void UdpServer::start(size_t numThreads)
{
    assert(numThreads > 0 && threads_.empty());
    threads_.reserve(numThreads);
    safeStartReceive();

    for (size_t i = 0; i != numThreads; ++i)
    {
        threads_.emplace_back(new clusnet::base::Thread([this]{ io_.run(); }));
        (*threads_.rbegin())->start();
    }

    LOG(INFO) << "udp server start ok. " << numThreads << " threads and 1 io_context";
}

void UdpServer::safeStartReceive()
{
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this]()
    {
        if (stop_)
        {
            return;
        }
        socket_.async_receive_from(boost::asio::buffer(recvBuffer_), remoteEndpoint_, boost::asio::bind_executor(strand_, std::bind(&UdpServer::handleReceive, this, std::placeholders::_1, std::placeholders::_2)));
    }));
}

void UdpServer::handleReceive(boost::system::error_code const& ec, std::size_t length)
{
    if (!ec)
    {
        auto it = sessions_.find(remoteEndpoint_);
        if (it == sessions_.end())
        {
            auto result = sessions_.emplace(remoteEndpoint_, std::make_shared<UdpSession>(io_, strand_, socket_, remoteEndpoint_, writeCallback_));
            assert(result.second == true);
            it = result.first;
        }
        assert(messageCallback_ != nullptr);
        UdpUniqDataPacketPtr msg(new UdpDataPacket(length));
        msg->append(recvBuffer_.data(), length);
        messageCallback_(it->second, std::move(msg));
    }
    else
    {
        LOG(WARNING) << "receive failed. " << ec.message();
    }
    safeStartReceive();
}

void UdpServer::stop()
{
    boost::asio::dispatch(io_, boost::asio::bind_executor(strand_, [this]
    {
        stop_ = true;
        socket_.close();
    }));

    std::for_each(threads_.begin(), threads_.end(), std::bind(&clusnet::base::Thread::join, std::placeholders::_1));
    threads_.clear();
    sessions_.clear();
}
